// SHUtility.cs
// Extension methods for sampling SphericalHarmonicsL2

using UnityEngine;
using UnityEngine.Rendering;

public static class SHUtility
{
    /// <summary>
    /// Sample a single direction from SphericalHarmonicsL2 and return the resulting color.
    /// Uses Unity's static Evaluate method under the hood.
    /// </summary>
    /// <param name="sh">The SphericalHarmonicsL2 instance.</param>
    /// <param name="direction">World-space direction to sample.</param>
    /// <returns>RGB color result of SH evaluation.</returns>
    public static Color EvaluateSHL2(this SphericalHarmonicsL2 sh, Vector3 direction)
    {
        // Normalize direction to ensure correct SH basis evaluation
        direction.Normalize();

        // Evaluate static method requires arrays
        Vector3[] dirs = new Vector3[1] { direction };
        Color[] results = new Color[1];

        // Perform evaluation
        sh.Evaluate(dirs, results);

        // Return the first (and only) result
        return results[0];
    }

    // just L0
    public static Color EvaluateL0(this SphericalHarmonicsL2 sh)
    {
        var coeffs = GetShaderConstantsFromNormalizedSH(sh);
        return new Color(coeffs[0].w, coeffs[1].w, coeffs[2].w);
    }
    
    public static Vector3 EvaluateL0_Vector3(this SphericalHarmonicsL2 sh)
    {
        var coeffs = GetShaderConstantsFromNormalizedSH(sh);
        return new Vector3(coeffs[0].w, coeffs[1].w, coeffs[2].w);
        //return new Vector3(sh[0,4], sh[1,4], sh[2,4]);
    }
    
    // just L1
    public static Color EvaluateL1(this SphericalHarmonicsL2 sh, Vector3 dir)
    {
        var coeffs = GetShaderConstantsFromNormalizedSH(sh);
        Vector4 dir4 = new Vector4(dir.x, dir.y, dir.z, 0f); // 忽略 L0 分量

        return new Color(
            Vector4.Dot(coeffs[0], dir4),
            Vector4.Dot(coeffs[1], dir4),
            Vector4.Dot(coeffs[2], dir4)
        );
    }
    
    public static Vector3 EvaluateL1_Vector3(this SphericalHarmonicsL2 sh, Vector3 dir)
    {
        var coeffs = GetShaderConstantsFromNormalizedSH(sh);
        Vector4 dir4 = new Vector4(dir.x, dir.y, dir.z, 0f); // 忽略 L0 分量

        return new Vector3(
            Vector4.Dot(coeffs[0], dir4),
            Vector4.Dot(coeffs[1], dir4),
            Vector4.Dot(coeffs[2], dir4)
        );
    }
    
    private static float MaxAbsChannelValue(Vector3 Input)
    {
        float MinValue = Mathf.Min(Input.x, Input.y, Input.z);
        float MaxValue = Mathf.Max(Input.x, Input.y, Input.z);
        
        return Mathf.Abs(MinValue) > Mathf.Abs(MaxValue) ? MinValue : MaxValue;
    }
    
    public static Vector3 EvaluateL1Luminance(this SphericalHarmonicsL2 sh)
    {
        //var coefficients = GetShaderConstantsFromNormalizedSH(sh);
        
        Vector4[] coefficients = new Vector4[3];
        for (int j = 0; j < 3; j++)
        {
            coefficients[j].x = sh[j, 3];
            coefficients[j].y = sh[j, 1];
            coefficients[j].z = sh[j, 2];
            coefficients[j].w = sh[j, 0] - sh[j, 6];
        }
        
        Vector3 SH1Lum = Vector3.zero;
        for (int j = 0; j < 3; j++)
        {
            Vector3 SH1Color;
            SH1Color.x = coefficients[0][j];
            SH1Color.y = coefficients[1][j];
            SH1Color.z = coefficients[2][j];

            SH1Color = 0.5f * (SH1Color + Vector3.one);

            float maxRGB = Mathf.Max(SH1Color.x, SH1Color.y, SH1Color.z);

            SH1Lum[j] = maxRGB;
        }

        return SH1Lum;
    }
    
    public static Vector3 EvaluateL1Luminance_1(this SphericalHarmonicsL2 sh)
    {
        var coeffs = GetShaderConstantsFromNormalizedSH(sh);
       
        Vector3 SH1Lum = Vector3.zero;
        for (int j = 0; j < 3; j++)
        {
            Vector3 SH1Color;
            SH1Color.x = coeffs[0][j]; //sh[0,j];
            SH1Color.y = coeffs[1][j]; //sh[1,j];
            SH1Color.z = coeffs[2][j]; //sh[2,j];

            SH1Color = 0.5f * (SH1Color + Vector3.one);

            float maxRGB = Mathf.Max(SH1Color.x, SH1Color.y, SH1Color.z);

            SH1Lum[j] = maxRGB;
        }

        return SH1Lum;
    }

    // just L2
    public static Color EvaluateL2(this SphericalHarmonicsL2 sh, Vector3 dir)
    {
        var coeffs = GetShaderConstantsFromNormalizedSH(sh);

        // L2 Part 1: xy, yz, zz, zx
        Vector4 vB = new Vector4(
            dir.x * dir.y,
            dir.y * dir.z,
            dir.z * dir.z,
            dir.z * dir.x
        );

        Vector3 part1 = new Vector3(
            Vector4.Dot(coeffs[3], vB),
            Vector4.Dot(coeffs[4], vB),
            Vector4.Dot(coeffs[5], vB)
        );

        // L2 Part 2: x² - y²
        float vC = dir.x * dir.x - dir.y * dir.y;

        Vector3 part2 = new Vector3(
            coeffs[6].x * vC,
            coeffs[6].y * vC,
            coeffs[6].z * vC
        );

        return new Color(part1.x + part2.x,part1.y + part2.y, part1.z + part2.z);
    }
    
    public static Color EvaluateL0L1(this SphericalHarmonicsL2 sh, Vector3 direction)
    {
        Vector4[] coeffs = GetShaderConstantsFromNormalizedSH(sh);

        Vector4 dir4 = new Vector4(direction.x, direction.y, direction.z, 1f); // 包含 L0 分量

        return new Color(
            Vector4.Dot(coeffs[0], dir4),
            Vector4.Dot(coeffs[1], dir4),
            Vector4.Dot(coeffs[2], dir4)
        );
    }
    
    public static Vector3 EvaluateL0L1_Vector3(this SphericalHarmonicsL2 sh, Vector3 direction)
    {
        Vector4[] coeffs = GetShaderConstantsFromNormalizedSH(sh);

        Vector4 dir4 = new Vector4(direction.x, direction.y, direction.z, 1f); // 包含 L0 分量

        return new Vector3(
            Vector4.Dot(coeffs[0], dir4),
            Vector4.Dot(coeffs[1], dir4),
            Vector4.Dot(coeffs[2], dir4)
        );
    }
    
    public static Color EvaluateFull(this SphericalHarmonicsL2 sh, Vector3 direction)
    {
        Vector4[] simdCoeffs = GetShaderConstantsFromNormalizedSH(sh);
        //return new Color(sh[0,0],sh[0,0],sh[0,0],sh[0,0]);
        Vector4 dir4 = new Vector4(direction.x, direction.y, direction.z, 1f);
        
        // L0 + L1
        Vector3 linearAndConstTerm = new Vector3(
            Vector4.Dot(simdCoeffs[0], dir4),
            Vector4.Dot(simdCoeffs[1], dir4),
            Vector4.Dot(simdCoeffs[2], dir4)
        );

        // L2 - part1 (xy, yz, zz, zx)
        Vector4 vB = new Vector4(
            direction.x * direction.y,
            direction.y * direction.z,
            direction.z * direction.z,
            direction.z * direction.x
        );

        Vector3 quadratic1 = new Vector3(
            Vector4.Dot(simdCoeffs[3], vB),
            Vector4.Dot(simdCoeffs[4], vB),
            Vector4.Dot(simdCoeffs[5], vB)
        );

        // L2 - part2 (x² - y²)
        float vC = direction.x * direction.x - direction.y * direction.y;
        Vector3 quadratic2 = new Vector3(
            simdCoeffs[6].x * vC,
            simdCoeffs[6].y * vC,
            simdCoeffs[6].z * vC
        );

        Vector3 final = linearAndConstTerm + quadratic1 + quadratic2;
        
        return new Color(final.x, final.y, final.z, 1f);
    }
    
    public static Vector4 EvaluateFullV(this SphericalHarmonicsL2 sh, Vector3 direction)
    {
        Vector4[] simdCoeffs = GetShaderConstantsFromNormalizedSH(sh);
        //return new Vector4(sh[0,0],sh[0,0],sh[0,0],sh[0,0]);
        Vector4 dir4 = new Vector4(direction.x, direction.y, direction.z, 1f);
        
        // L0 + L1
        Vector3 linearAndConstTerm = new Vector3(
            Vector4.Dot(simdCoeffs[0], dir4),
            Vector4.Dot(simdCoeffs[1], dir4),
            Vector4.Dot(simdCoeffs[2], dir4)
        );

        // L2 - part1 (xy, yz, zz, zx)
        Vector4 vB = new Vector4(
            direction.x * direction.y,
            direction.y * direction.z,
            direction.z * direction.z,
            direction.z * direction.x
        );

        Vector3 quadratic1 = new Vector3(
            Vector4.Dot(simdCoeffs[3], vB),
            Vector4.Dot(simdCoeffs[4], vB),
            Vector4.Dot(simdCoeffs[5], vB)
        );

        // L2 - part2 (x² - y²)
        float vC = direction.x * direction.x - direction.y * direction.y;
        Vector3 quadratic2 = new Vector3(
            simdCoeffs[6].x * vC,
            simdCoeffs[6].y * vC,
            simdCoeffs[6].z * vC
        );

        Vector3 final = linearAndConstTerm + quadratic1 + quadratic2;
        
        return new Vector4(final.x, final.y, final.z, 1f);
    }
    
    /// <summary>
    /// 模拟 Unity 的 SphericalHarmonicsL2::GetShaderConstantsFromNormalizedSH
    /// 将 SH 系数重新打包为 7 个 float4。
    /// </summary>
    public static Vector4[] GetShaderConstantsFromNormalizedSH(SphericalHarmonicsL2 sh)
    {
        Vector4[] outCoefficients = new Vector4[7];

        // R/G/B -> 0/1/2
        for (int iC = 0; iC < 3; iC++)
        {
            // Swizzle the coefficients to be in { x, y, z, DC } order.
            float x = sh[iC, 3]; // L1.y
            float y = sh[iC, 1]; // L1.x
            float z = sh[iC, 2]; // L1.z
            float w = sh[iC, 0] - sh[iC, 6]; // L0 - L2.2

            outCoefficients[iC] = new Vector4(x, y, z, w);
        }

        for (int iC = 0; iC < 3; iC++)
        {
            float x = sh[iC, 4];        // L2.0
            float y = sh[iC, 5];        // L2.1
            float z = sh[iC, 6] * 3.0f; // L2.2 * 3
            float w = sh[iC, 7];        // L2.3

            outCoefficients[iC + 3] = new Vector4(x, y, z, w);
        }

        // Final L2.4 term (coefficient 8) across R, G, B
        float r = sh[0, 8]; // red
        float g = sh[1, 8]; // green
        float b = sh[2, 8]; // blue

        outCoefficients[6] = new Vector4(r, g, b, 1.0f); // last slot

        return outCoefficients;
    }

}